{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Haruhi-2-Dev/blob/main/notebook/try_chroma.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Welcome to this very small demo of Chroma\n",
        "\n",
        "1. Install chroma\n",
        "2. Create the client\n",
        "3. Create the collection\n",
        "4. Load data\n",
        "5. Query data by nearest embedding"
      ],
      "metadata": {
        "id": "eLtiBx5FOWer"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "jbwLWGiDElhM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3d132a2a-b9f6-4870-81c4-d205946de7aa"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m402.8/402.8 kB\u001b[0m \u001b[31m10.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m78.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.4/58.4 kB\u001b[0m \u001b[31m5.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.5/59.5 kB\u001b[0m \u001b[31m6.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.3/5.3 MB\u001b[0m \u001b[31m16.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.9/5.9 MB\u001b[0m \u001b[31m120.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m126.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.3/67.3 kB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.0/67.0 kB\u001b[0m \u001b[31m5.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m4.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m428.8/428.8 kB\u001b[0m \u001b[31m38.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.1/4.1 MB\u001b[0m \u001b[31m112.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m79.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m129.9/129.9 kB\u001b[0m \u001b[31m14.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m9.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for chroma-hnswlib (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for pypika (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "pip -q install chromadb"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import chromadb\n",
        "client = chromadb.PersistentClient(path = '/content/mama')\n",
        "\n",
        "collection = client.get_or_create_collection(\"test\")\n",
        "\n"
      ],
      "metadata": {
        "id": "ciUVixVxwiws"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "collection.add(\n",
        "    embeddings=[\n",
        "        [1.1, 2.3, 3.2],\n",
        "        [4.5, 6.9, 4.4],\n",
        "        [1.1, 2.3, 3.2],\n",
        "        [4.5, 6.9, 4.4],\n",
        "        [1.1, 2.3, 3.2],\n",
        "        [4.5, 6.9, 4.4],\n",
        "        [1.1, 2.3, 3.2],\n",
        "        [4.5, 6.9, 4.4],\n",
        "    ],\n",
        "    metadatas=[\n",
        "        {\"uri\": \"img1.png\", \"style\": \"style1\"},\n",
        "        {\"uri\": \"img2.png\", \"style\": \"style2\"},\n",
        "        {\"uri\": \"img3.png\", \"style\": \"style1\"},\n",
        "        {\"uri\": \"img4.png\", \"style\": \"style1\"},\n",
        "        {\"uri\": \"img5.png\", \"style\": \"style1\"},\n",
        "        {\"uri\": \"img6.png\", \"style\": \"style1\"},\n",
        "        {\"uri\": \"img7.png\", \"style\": \"style1\"},\n",
        "        {\"uri\": \"img8.png\", \"style\": \"style1\"},\n",
        "    ],\n",
        "    documents=[\"doc1\", \"doc2\", \"doc3\", \"doc4\", \"doc5\", \"doc6\", \"doc7\", \"doc8\"],\n",
        "    ids=[\"id1\", \"id2\", \"id3\", \"id4\", \"id5\", \"id6\", \"id7\", \"id8\"],\n",
        ")\n",
        "\n",
        "query_result = collection.query(\n",
        "        query_embeddings=[[1.1, 2.3, 3.2], [5.1, 4.3, 2.2]],\n",
        "        n_results=2,\n",
        "    )\n",
        "\n",
        "print(query_result)\n"
      ],
      "metadata": {
        "id": "64dWyeEdKAX9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5c4a64a9-5ed3-4bb2-f190-f33829492bb5"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'ids': [['id1', 'id3'], ['id2', 'id8']], 'distances': [[0.0, 0.0], [11.960000000000003, 11.960000000000003]], 'metadatas': [[{'style': 'style1', 'uri': 'img1.png'}, {'style': 'style1', 'uri': 'img3.png'}], [{'style': 'style2', 'uri': 'img2.png'}, {'style': 'style1', 'uri': 'img8.png'}]], 'embeddings': None, 'documents': [['doc1', 'doc3'], ['doc2', 'doc8']]}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "embeddings=[\n",
        "        [1.1, 2.3, 3.6],\n",
        "        [4.5, 6.9, 4.4],\n",
        "        [1.1, 2.3, 3.2],\n",
        "        [4.5, 6.9, 4.4],\n",
        "        [1.1, 2.3, 3.3],\n",
        "        [4.5, 6.9, 4.4],\n",
        "        [1.1, 2.3, 3.2],\n",
        "        [4.5, 6.9, 4.4],\n",
        "    ]"
      ],
      "metadata": {
        "id": "lXeEEdL-NUj9"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "documents=[\"doc1\", \"doc2\", \"doc3\", \"doc4\", \"doc5\", \"doc6\", \"doc7\", \"doc8\"]"
      ],
      "metadata": {
        "id": "mLC2jeyWOqjS"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import chromadb\n",
        "# from BaseDB import BaseDB\n",
        "import random\n",
        "import string\n",
        "import os\n",
        "\n",
        "class ChromaDB:\n",
        "\n",
        "    def __init__(self):\n",
        "        self.client = None\n",
        "        self.collection = None\n",
        "        self.path = None\n",
        "\n",
        "    def init_db(self):\n",
        "\n",
        "        folder_name = ''\n",
        "\n",
        "        while os.path.exists(folder_name) or folder_name == '':\n",
        "            # try to create a folder named temp_<random string> which is not yet existed\n",
        "            folder_name =  \"tempdb_\" + ''.join(random.sample(string.ascii_letters + string.digits, 8))\n",
        "\n",
        "        self.path = folder_name\n",
        "        self.client = chromadb.PersistentClient(path = folder_name)\n",
        "\n",
        "        self.collection = self.client.get_or_create_collection(\"search\")\n",
        "\n",
        "    def save(self, file_path):\n",
        "        if file_path != self.path:\n",
        "            # copy all files in self.path to file_path\n",
        "            os.system(\"cp -r {} {}\".format(self.path, file_path))\n",
        "            self.path = file_path\n",
        "\n",
        "\n",
        "    def load(self, file_path):\n",
        "        self.path = file_path\n",
        "        self.client = chromadb.PersistentClient(path = file_path)\n",
        "        self.collection = self.client.get_collection(\"search\")\n",
        "\n",
        "    def search(self, vector, n_results):\n",
        "        results = self.collection.query(query_embeddings=[vector], n_results=n_results)\n",
        "        return results['documents'], results['distances']\n",
        "\n",
        "    def init_from_docs(self, vectors, documents):\n",
        "        if self.client is None:\n",
        "            self.init_db()\n",
        "        ids = []\n",
        "        for i, doc in enumerate(documents):\n",
        "            first_four_chat = doc[:min(4, len(doc))]\n",
        "            ids.append( str(i) + \"_\" + doc)\n",
        "        self.collection.add(embeddings=vectors, documents=documents, ids = ids)\n",
        "\n"
      ],
      "metadata": {
        "id": "s05ZVFAOOuAp"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "What's next?\n",
        "- [Read the docs](https://docs.trychroma.com/)\n",
        "- [Join the Discord](https://discord.gg/MMeYNTmh3x)\n",
        "- [View project on Github](https://github.com/chroma-core/chroma)"
      ],
      "metadata": {
        "id": "3o94AN6aOlgd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "firstdb = ChromaDB()\n",
        "firstdb.init_db()\n",
        "firstdb.init_from_docs(embeddings, documents )\n",
        "firstdb.save('/content/baba')\n",
        "print(firstdb.search( [1.1, 2.3, 3.2],4 ))"
      ],
      "metadata": {
        "id": "ExiIuwlsOyq6",
        "outputId": "74d892c9-4c5b-41e7-9196-514f9652297e",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "([['doc3', 'doc7', 'doc5', 'doc1']], [[0.0, 0.0, 0.009999999999999929, 0.15999999999999992]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "seconddb = ChromaDB()\n",
        "seconddb.load('/content/baba')\n",
        "print(firstdb.search( [1.1, 2.3, 3.2],4 ))"
      ],
      "metadata": {
        "id": "wiJX3rI0PDT_",
        "outputId": "b28bb9b0-5564-43e2-8aec-07c2a9f64495",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "([['doc3', 'doc7', 'doc5', 'doc1']], [[0.0, 0.0, 0.009999999999999929, 0.15999999999999992]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "4eHibkRaPfZM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}